import grpc
import calligraphy_pb2_grpc,calligraphy_pb2

import os
from PIL import Image,ImageDraw,ImageFont
import torchvision.transforms as transforms
import glob
import matplotlib.pyplot as plt
import numpy as np
import time
import torch
import base64
from torchvision.models.detection import fasterrcnn_resnet50_fpn
import itertools
import cv2

preprocessing_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225]),
])

preprocessing_val = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])


def run(image, height, width):
  with grpc.insecure_channel('localhost:2534') as channel:
  # with grpc.insecure_channel('207.246.117.29:8080') as channel:
    stub = calligraphy_pb2_grpc.NetStub(channel=channel)
    request = calligraphy_pb2.ImageMatrix(image=image, height=height, width=width)

    start = time.time()
    response = stub.Detection(request)
    end = time.time()

    bboxes = np.frombuffer(base64.b64decode(response.bboxes),dtype=np.float32)
    labels = np.frombuffer(base64.b64decode(response.labels),dtype=np.int)
    scores = np.frombuffer(base64.b64decode(response.scores),dtype=np.float32)


    run_time = response.run_time
    total_time = end - start
    image_transfer_time = total_time - run_time

    print("total time = ", end-start,
          "  forward time = ",run_time,
          "  image_transfer_time = ",image_transfer_time)
    return bboxes,labels,scores

if __name__ == '__main__':

  for path in glob.glob("/media/retoo/RetooDisk1/wanghui/Data/coco2014/test2014/*.jpg"):
    src = Image.open(path).convert('RGB')
    name = os.path.basename(path).strip()
    subdir = os.path.split(os.path.split(path)[0])[1]
    width, height = src.size
    image = np.asarray(src)

    ## encode image
    # str_encode = base64.b64encode(image)
    # _, ext = os.path.splitext(name)
    # img_data = cv2.imread(path)[:, :, ::-1]  # read image and convert to RGB channel
    # _, img_encode = cv2.imencode(ext, img_data)
    # data_encode = np.array(img_encode)
    # str_encode = data_encode.tobytes()

    ## decode image
    # nparr = np.frombuffer(str_encode, np.uint8)
    # decode_img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    encode_image = base64.b64encode(open(path, "rb").read())
    boxes, labels, scores = run(encode_image, height, width)

    # save detection result
    save_image = src.copy()
    draw = ImageDraw.Draw(save_image)
    colors = itertools.cycle(["red", "green", "blue", "yellow", "cyan", "gold", "purple", "violet", "pink"])
    boxes = np.reshape(boxes, [-1,4])
    for idx in range(boxes.shape[0]):
        xmin, ymin, xmax, ymax = boxes[idx]
        score = scores[idx]
        label = labels[idx]

        if score > 0.4:
            # crop_character = src.crop((xmin, ymin, xmax, ymax))
            # character_save_path = "crop_character/%s/" % subdir
            # os.makedirs(character_save_path, exist_ok=True)
            # crop_character.save("%s/%s_%d.jpg" % (character_save_path, name[:-4], idx))
            draw.rectangle(((xmin, ymin), (xmax, ymax)), fill=None, outline=next(colors), width=2)  ##绘制矩形框，指定外轮廓颜色和填充颜色
            draw.text((xmin, ymin - 10),
                      str(round(score, 2)),
                      font=ImageFont.truetype("simhei.ttf", 10),
                      fill=(255, 0, 0))  ##指定字体和颜色，需要复制字体到本地路径下

    print("process image : {}".format(path))
    save_path = "detection_result/%s" % subdir
    os.makedirs(save_path, exist_ok=True)
    save_image.save("%s/%s" % (save_path, name))

    del draw, src, save_image


# docker build --pull -t resnet-libtorch-serving -f libtorch_cpu_Dockerfile .
# docker run -p 45.77.51.15:8080:8080 --name=pytorch_service -d -it dl-pytorch-serving:latest python3 ./server.py
# docker run -p 207.246.117.29:8080:8080 --name=pytorch_service -it dl-pytorch-serving:latest python3 ./server.py

